import json
from collections import defaultdict
from django.http import JsonResponse
import cv2
import base64
import numpy as np
import time
import asyncio
from channels.generic.websocket import AsyncWebsocketConsumer
from ultralytics import YOLO
from django.utils import timezone
from channels.db import database_sync_to_async
from rest_framework.authtoken.models import Token
from detection.models import RealtimeDetectionResult, RealtimeMonitoring
import pytz
from datetime import datetime
import redis
from django.conf import settings
import pickle
from detection.tasks import process_session_data
import logging

# 全局加载 YOLO 模型
model = YOLO('yolov8n.pt')

# Redis连接
redis_client = redis.Redis(
    host=settings.CACHES['default']['LOCATION'].split('//')[1].split(':')[0],
    port=int(settings.CACHES['default']['LOCATION'].split(':')[2].split('/')[0]),
    db=1,  # 使用db1存储实时检测数据
    decode_responses=False  # 不自动解码，因为我们要存储二进制数据
)

logger = logging.getLogger(__name__)

class YoloDetectionConsumer(AsyncWebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 用户相关
        self.user = None
        self.user_id = None
        
        # 会话相关
        self.detection_session = None
        self.session_active = False
        self.session_timeout = 30.0  # 会话超时时间（秒）
        
        # 目标检测相关
        self.active_objects = {}  # 当前活跃的目标
        self.object_timeout = 0.5  # 目标超时时间（秒）
        
        # 清理相关
        self.last_cleanup_time = time.time()
        self.cleanup_interval = 1.0  # 清理间隔（秒）
        
        # 帧处理相关
        self.frame_start_time = None
        self.last_frame_time = None
        self.inactivity_task = None
        
        # 数据处理相关
        self.last_process_time = None
        self.process_interval = 300  # 每5分钟处理一次数据
        
        # 时区设置
        self.timezone = pytz.timezone('Asia/Shanghai')

    async def connect(self):
        """处理WebSocket连接"""
        token = self.scope.get('query_string', b'').decode('utf-8').split('=')[-1]
        user = await self.get_user_from_token(token)
        
        if user is None:
            await self.close(code=4000)
            return
            
        self.user = user
        self.user_id = user.id
        
        # 创建新的检测会话
        self.detection_session = await self.create_detection_session()
        self.session_active = True
        self.inactivity_task = asyncio.create_task(self.monitor_inactivity())
        
        await self.accept()

    async def disconnect(self, close_code):
        """处理WebSocket断开连接"""
        try:
            if self.session_active and self.detection_session:
                # 确保在断开连接时处理所有数据
                await database_sync_to_async(process_session_data)(str(self.detection_session.session_id), is_final=True)
                # 更新会话结束时间
                self.detection_session.end_time = timezone.now()
                await database_sync_to_async(self.detection_session.save)()
                self.session_active = False
                logger.info(f"WebSocket连接已关闭，会话 {self.detection_session.session_id} 数据已保存")
            
            if self.inactivity_task and not self.inactivity_task.done():
                self.inactivity_task.cancel()
        except Exception as e:
            logger.error(f"处理会话数据时出错: {str(e)}", exc_info=True)
        finally:
            # 确保清理所有资源
            self.session_active = False
            self.detection_session = None
        logger.info(f"WebSocket连接已关闭: {close_code}")

    @database_sync_to_async
    def get_user_from_token(self, token):
        """通过Token获取用户"""
        try:
            token_obj = Token.objects.get(key=token)
            return token_obj.user
        except Token.DoesNotExist:
            return None

    def get_current_time(self):
        """获取当前时间（带时区）"""
        logger.info(f"获取当前时间: {timezone.now()}")
        return timezone.now()

    @database_sync_to_async
    def create_detection_session(self):
        """创建新的检测会话"""
        if not self.session_active:
            current_time = timezone.now()
            session = RealtimeMonitoring.objects.create(
                user_id=self.user_id,
                start_time=current_time
            )
            self.session_active = True
            self.detection_session = session
            self.last_process_time = time.time()
            return session
        return self.detection_session

    @database_sync_to_async
    def close_detection_session(self):
        """结束当前检测会话"""
        if self.detection_session and self.session_active:
            self.detection_session.end_time = timezone.now()
            self.detection_session.save()
            self.session_active = False
            
            # 立即处理会话数据
            process_session_data(self.detection_session.id)

    async def monitor_inactivity(self):
        """监控会话活跃状态"""
        try:
            while self.session_active:
                await asyncio.sleep(1)  # 每秒检查一次
                current_time = time.time()
                
                # 检查会话是否超时
                if self.last_frame_time and (current_time - self.last_frame_time > self.session_timeout):
                    print(f"Session timeout after {self.session_timeout} seconds of inactivity")
                    await self.close_detection_session()
                    break
                
                # 检查是否需要处理数据（每5分钟）
                if self.last_process_time and (current_time - self.last_process_time > self.process_interval):
                    print(f"Processing data after {self.process_interval} seconds")
                    # 处理当前会话的数据
                    process_session_data(self.detection_session.id)
                    self.last_process_time = current_time
                    
        except asyncio.CancelledError:
            pass

    def save_frame_to_redis(self, frame_data):
        """将帧数据保存到Redis，使用hash结构"""
        if not self.detection_session:
            return

        try:
            # 构造Redis键，使用主键id
            session_key = f"detection:session:{self.detection_session.id}"
            frame_id = f"{int(time.time() * 1000)}"  # 使用毫秒时间戳作为帧ID

            # 确保frame_data包含所有必要的数据
            if not isinstance(frame_data, dict):
                logger.error(f"Error: frame_data is not a dict: {type(frame_data)}")
                return

            if 'detections' not in frame_data:
                logger.error(f"Error: no detections in frame_data: {frame_data}")
                return

            # 序列化数据
            frame_data['session_id'] = self.detection_session.id
            frame_data['timestamp'] = time.time()
            
            serialized_data = pickle.dumps(frame_data)

            # 使用HSET保存到Redis hash
            redis_client.hset(session_key, frame_id, serialized_data)
            redis_client.expire(session_key, 3600)  # 设置1小时过期时间

            logger.debug(f"Saved frame {frame_id} to Redis with key {session_key}")

        except Exception as e:
            logger.error(f"Error saving frame to Redis: {str(e)}", exc_info=True)

    async def receive(self, text_data=None, bytes_data=None):
        """处理接收到的帧数据"""
        if not text_data:
            return
            
        try:
            # 更新帧时间戳
            self.frame_start_time = time.time()
            self.last_frame_time = time.time()

            # 解析帧数据
            data = json.loads(text_data)
            img_data = data.get('frame')
            if not img_data:
                await self.send(text_data=json.dumps({'error': 'No frame data'}))
                return

            # 处理帧
            result = await self.process_frame(img_data)
            
            # 保存帧数据到Redis
            self.save_frame_to_redis(result)
            
            # 发送结果给客户端
            await self.send(text_data=json.dumps(result))

        except Exception as e:
            print(f"Error processing frame: {str(e)}")
            await self.send(text_data=json.dumps({'error': str(e)}))

    @database_sync_to_async
    def process_frame(self, img_data):
        """处理图像帧并执行目标检测"""
        try:
            # 解码图像数据
            if ',' in img_data:
                img_data = img_data.split(',')[1]
            img_bytes = base64.b64decode(img_data)
            img_np = np.frombuffer(img_bytes, np.uint8)
            frame = cv2.imdecode(img_np, cv2.IMREAD_COLOR)

            if frame is None:
                return {'error': 'Could not decode image'}

            return self._process_detection(frame)

        except Exception as e:
            print(f"Frame processing error: {str(e)}")
            return {'error': str(e)}

    def _process_detection(self, frame):
        """执行YOLO检测并处理结果"""
        current_time = time.time()
        results = model(frame, verbose=False)
        detections = []
        current_detections = set()

        for r in results:
            boxes = r.boxes
            for i, box in enumerate(boxes):
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                confidence = float(box.conf[0])
                class_id = int(box.cls[0])
                class_name = model.names[class_id]

                if confidence > 0.5:
                    # 生成目标唯一标识
                    object_id = f"{class_id}_{int(x1)}_{int(y1)}_{int(x2)}_{int(y2)}"
                    current_detections.add(object_id)

                    # 更新或创建目标记录
                    if object_id not in self.active_objects:
                        self.active_objects[object_id] = {
                            'class_id': class_id,
                            'class_name': class_name,
                            'start_time': self.frame_start_time,
                            'last_seen': current_time,
                            'confidence': confidence
                        }
                    else:
                        self.active_objects[object_id].update({
                            'last_seen': current_time,
                            'confidence': max(self.active_objects[object_id]['confidence'], confidence)
                        })

                    detections.append({
                        'id': i,
                        'bbox': [x1, y1, x2, y2],
                        'confidence': confidence,
                        'class_id': class_id,
                        'class_name': class_name,
                        'timestamp': current_time
                    })

        return {
            'success': True,
            'detections': detections,
            'frame_timestamp': current_time
        }

    @database_sync_to_async
    def cleanup_objects(self, force_cleanup=False):
        """清理超时目标并保存检测结果"""
        current_time = time.time()

        # 检查是否需要执行清理
        if not force_cleanup and (current_time - self.last_cleanup_time < self.cleanup_interval):
            return

        # 处理超时目标
        objects_to_remove = []
        for object_id, obj in self.active_objects.items():
            if force_cleanup or (current_time - obj['last_seen'] > self.object_timeout):
                if self.detection_session and self.detection_session.id:
                    # 将时间戳转换为带时区的datetime对象
                    start_time = datetime.fromtimestamp(obj['start_time'], tz=self.timezone)
                    end_time = datetime.fromtimestamp(obj['last_seen'], tz=self.timezone)
                    
                    RealtimeDetectionResult.objects.create(
                        session_id=self.detection_session.id,
                        class_id=obj['class_id'],
                        class_name=obj['class_name'],
                        start_time=start_time.timestamp(),  # 存储为时间戳
                        end_time=end_time.timestamp(),      # 存储为时间戳
                        confidence=obj['confidence']
                    )
                objects_to_remove.append(object_id)

        # 移除已处理的目标
        for object_id in objects_to_remove:
            del self.active_objects[object_id]

        self.last_cleanup_time = current_time